package com.atguigu.dga.governance.assessor.calc;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.atguigu.dga.constant.DgaConstant;
import com.atguigu.dga.governance.assessor.Assessor;
import com.atguigu.dga.governance.bean.AssessParam;
import com.atguigu.dga.governance.bean.GovernanceAssessDetail;
import com.atguigu.dga.meta.bean.TableMetaInfo;
import com.atguigu.dga.util.SqlUtil;
import com.google.common.collect.Sets;
import lombok.Getter;
import org.apache.commons.collections.CollectionUtils;
import org.apache.hadoop.hive.ql.lib.Dispatcher;
import org.apache.hadoop.hive.ql.lib.Node;
import org.apache.hadoop.hive.ql.parse.ASTNode;
import org.apache.hadoop.hive.ql.parse.HiveParser;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.math.BigDecimal;
import java.net.URISyntaxException;
import java.text.ParseException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * ClassName: SqlSimpleProcessAssessor
 * Package: com.atguigu.dga.governance.assessor.calc
 * Description:
 *
 * @Author:
 * @Create: 2024/3/20 15:09
 * @Version: 1.0
 */
@Component("SQL_SIMPLE_PROCESS")
public class SqlSimpleProcessAssessor extends Assessor {
    @Value("${default.datawarehouse.database}")
    private String defaultDatabase;

    @Override
    public void checkProblem(AssessParam assessParam, GovernanceAssessDetail governanceAssessDetail) throws ParseException, URISyntaxException, IOException, InterruptedException, org.apache.hadoop.hive.ql.parse.ParseException, SemanticException {
        // sql中提取 复杂计算 Join union groupby 和分区字段
        // 判断是否有复杂计算 过滤字段是非分区字段

        // TODO 1 ods 表没有sql
        if(DgaConstant.DW_LEVEL_ODS.equals(assessParam.getTableMetaInfo().getTableMetaInfoExtra().getDwLevel())) {
            return;
        }

        // TODO 1 排除其他没有sql的
        if(assessParam.getTDsTaskDefinition()== null || assessParam.getTDsTaskDefinition().getTaskSQl() == null) {
            return;
        }

        // TODO 2 取sql
        String taskSql = assessParam.getTDsTaskDefinition().getTaskSQl();

        // TODO 3 处理sql
        SimpleProcessDispatcher simpleProcessDispatcher = new SimpleProcessDispatcher(defaultDatabase);
        SqlUtil.parseSql(taskSql, simpleProcessDispatcher);

        // TODO 3-3 判断是否包含复杂计算
        Set<String> realComplicateSet = simpleProcessDispatcher.getRealComplicateSet();
        if( realComplicateSet.size() > 0) {
            // 存在复杂计算
            governanceAssessDetail.setAssessComment("存在的复杂计算：" + realComplicateSet);
        }

        // TODO 4 判断过滤字段是否为非分区字段
        Set<String> realOperateSet = simpleProcessDispatcher.realOperateSet;
        Set<String> realTableNameSet = simpleProcessDispatcher.realTableNameSet;

        // TODO 5 准备好所有的tableMetaInfoList 转为map结构
        List<TableMetaInfo> tableMetaInfoList = assessParam.getTableMetaInfoList();
        Map<String, TableMetaInfo> tableMetaInfoHashMap = new HashMap<>(tableMetaInfoList.size());
        for (TableMetaInfo tableMetaInfo : tableMetaInfoList) {
            tableMetaInfoHashMap.put(tableMetaInfo.getSchemaName() + "." + tableMetaInfo.getTableName(), tableMetaInfo);
        }

        // TODO 5 维护被查询表的分区字段 遍历from的表集合realTableNameSet
        Set<String> allPartitionNameSet = new HashSet<>();
        for (String tableName : realTableNameSet) {
            String partitionColNameJson = null;
            try {
                TableMetaInfo tableMetaInfo = tableMetaInfoHashMap.get(tableName);
                partitionColNameJson = tableMetaInfo.getPartitionColNameJson();
            } catch (Exception e) {
                System.out.println("tableName = " + tableName);  // TODO bug?????
                throw new RuntimeException(e);
            }
            List<JSONObject> partitionJsonObjList = JSON.parseArray(partitionColNameJson, JSONObject.class);
            // 取出分区字段名
            List<String> partitionNameList = partitionJsonObjList.stream()
                    .map(jsonObj -> jsonObj.getString("name")).collect(Collectors.toList());
            allPartitionNameSet.addAll(partitionNameList);
        }

        // TODO 6 集合作差 被查询表的分区字段 与 实际存在的过滤字段
        Collection notPartitionNameList = CollectionUtils.subtract(allPartitionNameSet, realOperateSet);

        // TODO 7 判断 如果存在差，setComment 把差（非分区字段的过滤）记录下来
        if(notPartitionNameList.size() > 0) {
            governanceAssessDetail.setAssessComment("存在的非分区过滤字段：" + notPartitionNameList);
        }

        // TODO 8 简单加工
        governanceAssessDetail.setAssessScore(BigDecimal.ZERO);
        governanceAssessDetail.setAssessProblem("SQL 为简单加工");
    }

    public static class SimpleProcessDispatcher implements Dispatcher {
        private String defaultDatabase;

        public SimpleProcessDispatcher(String defaultDatabase) {
            this.defaultDatabase = defaultDatabase;
        }

        //TODO 3-1 定义一个集合， 哪些计算是复杂计算
        Set<Integer> complicateTokSet = Sets.newHashSet(
                HiveParser.TOK_JOIN ,   // join ,包含通过where连接的情况
                HiveParser.TOK_GROUPBY , // group by
                HiveParser.TOK_LEFTOUTERJOIN , // left join
                HiveParser.TOK_RIGHTOUTERJOIN , //right join
                HiveParser.TOK_FULLOUTERJOIN , // full join
                HiveParser.TOK_FUNCTION , // count(1)
                HiveParser.TOK_FUNCTIONDI, // count(distinct xx)
                HiveParser.TOK_FUNCTIONSTAR , // count(*)
                HiveParser.TOK_SELECTDI , // distinct
                HiveParser.TOK_UNIONALL // union
        ) ;
        // TODO 3-2 定义一个集合，维护复杂计算
        @Getter
        Set<String> realComplicateSet = new HashSet<>();

        // TODO 4-1 定义集合， where后面条件的操作符号
        Set<Integer> operatorSet = Sets.newHashSet(
                HiveParser.EQUAL ,   // =
                HiveParser.GREATERTHAN , // >
                HiveParser.LESSTHAN, // <
                HiveParser.GREATERTHANOREQUALTO , // >=
                HiveParser.LESSTHANOREQUALTO , // <=
                HiveParser.NOTEQUAL , // <>
                HiveParser.KW_LIKE // like
        ) ;

        // TODO 4-1 定义集合 维护实际存在的过滤字段
        @Getter
        Set<String> realOperateSet = new HashSet<>();

        // TODO 4-2 定义集合 维护from的表
        @Getter
        Set<String> realTableNameSet = new HashSet<>();

        @Override
        public Object dispatch(Node node, Stack<Node> stack, Object... objects) throws SemanticException {
            ASTNode astNode = (ASTNode) node;
            // TODO 3-2 判断是否为复杂计算 保存下来
            if(complicateTokSet.contains(astNode.getType())) {
                realComplicateSet.add(astNode.getText());
            }

            // TODO 4-1 提取过滤字段
            // 找到 符号对应的节点，判断该节点的祖先是否有TOK_WHERE 如果有，该节点就是where后面的过滤
            if(operatorSet.contains(astNode.getType()) && astNode.getAncestor(HiveParser.TOK_WHERE) != null) {
                // 获取第一个孩子 可能是TOK_TABLE_OR_COL  也可能是DOT 有起别名
                ASTNode firstChild = (ASTNode)astNode.getChild(0);
                if(firstChild.getType() == HiveParser.TOK_TABLE_OR_COL) {
                    realOperateSet.add(firstChild.getChild(0).getText());
                } else if (firstChild.getType() == HiveParser.DOT) {
                    realOperateSet.add(firstChild.getChild(1).getText());
                }
            }

            // TODO 4-2 提取from的表  取表名有问题？？？？？？？？
//            if(astNode.getType() == HiveParser.TOK_TABNAME) {
//                ArrayList<Node> children = astNode.getChildren();
//                System.out.println("children = " + children);
//                if(children.size() == 1) {
//                    // 一个孩子 直接取 就是表名  统一一下，库名通过配置传过来
//                    realTableNameSet.add(defaultDatabase + "." + astNode.getChild(0).getText());
//                }else if(children.size() == 2) {
//                    // 两个孩子 第一个是库名 第二个是表名
//                    realTableNameSet.add(astNode.getChild(0).getText() + "." + astNode.getChild(1).getText());
//                }
//            }

            if (astNode.getType() == HiveParser.TOK_TABNAME) {
                ArrayList<Node> children = astNode.getChildren();
//                System.out.println("children = " + children);
                if (children.size() == 1) {
                    // 如果只有一个孩子，检查孩子节点类型是否为TOK_TABLE_OR_COL
                    ASTNode child = (ASTNode)children.get(0);
                    if (child.getType() == HiveParser.TOK_TABLE_OR_COL) {
                        realTableNameSet.add(defaultDatabase + "." + child.getChild(0).getText());
                    } else {
                        // 如果不是TOK_TABLE_OR_COL类型，则说明是表别名，根据实际情况处理
                        // 可以选择忽略或者进行其他操作
                    }
                } else if (children.size() == 2) {
                    // 两个孩子 第一个是库名 第二个是表名
                    realTableNameSet.add(astNode.getChild(0).getText() + "." + astNode.getChild(1).getText());
                }
            }


            return null;
        }
    }
}
